package whitelist

import (
	"fmt"
	"net"
	"sync"
	"time"

	"openSesame/pkg/logger"
)

// IPInfo 存储IP相关信息
type IPInfo struct {
	IsWhitelisted bool
	LastAccess    time.Time
}

// Whitelist 白名单存储
type Whitelist struct {
	ips        map[string]*IPInfo
	mu         sync.RWMutex
	expiration time.Duration
	stopChan   chan struct{}
	accessChan chan string
}

// New 创建新的白名单实例
func New(expirationSeconds int) *Whitelist {
	wl := &Whitelist{
		ips:        make(map[string]*IPInfo),
		expiration: time.Duration(expirationSeconds) * time.Second,
		stopChan:   make(chan struct{}),
		accessChan: make(chan string, 100), // 缓冲channel，避免阻塞
	}

	// 启动清理协程
	go wl.cleanupLoop()
	// 启动访问通知处理协程
	go wl.accessLoop()

	return wl
}

// Stop 停止白名单清理协程
func (w *Whitelist) Stop() {
	close(w.stopChan)
	close(w.accessChan)
}

// accessLoop 处理访问通知
func (w *Whitelist) accessLoop() {
	for {
		select {
		case ip, ok := <-w.accessChan:
			if !ok {
				return
			}
			w.updateAccessTime(ip)
		case <-w.stopChan:
			return
		}
	}
}

// updateAccessTime 更新IP访问时间
func (w *Whitelist) updateAccessTime(ip string) {
	w.mu.Lock()
	defer w.mu.Unlock()

	if info, exists := w.ips[ip]; exists {
		info.LastAccess = time.Now()
		logger.Debugf("Updated access time for IP: %s", ip)
	}
}

// NotifyAccess 通知IP访问
func (w *Whitelist) NotifyAccess(ip string) {
	select {
	case w.accessChan <- ip:
		// 成功发送通知
	default:
		// channel已满，记录警告日志
		logger.Warnf("Access notification channel is full, dropping notification for IP: %s", ip)
	}
}

// cleanupLoop 定期清理过期的IP
func (w *Whitelist) cleanupLoop() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			w.cleanup()
		case <-w.stopChan:
			return
		}
	}
}

// cleanup 清理过期的IP
func (w *Whitelist) cleanup() {
	w.mu.Lock()
	defer w.mu.Unlock()

	now := time.Now()
	for ip, info := range w.ips {
		if now.Sub(info.LastAccess) > w.expiration {
			delete(w.ips, ip)
			logger.Infof("IP %s removed from whitelist due to expiration", ip)
		}
	}
}

// Add 添加IP到白名单
func (w *Whitelist) Add(ip string) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.ips[ip] = &IPInfo{
		IsWhitelisted: true,
		LastAccess:    time.Now(),
	}
}

// Contains 检查IP是否在白名单中
func (w *Whitelist) Contains(ip string) bool {
	w.mu.RLock()
	defer w.mu.RUnlock()
	if info, exists := w.ips[ip]; exists {
		info.LastAccess = time.Now()
		return info.IsWhitelisted
	}
	return false
}

// GetLastAccess 获取IP最后访问时间
func (w *Whitelist) GetLastAccess(ip string) (time.Time, bool) {
	w.mu.RLock()
	defer w.mu.RUnlock()
	if info, exists := w.ips[ip]; exists {
		return info.LastAccess, true
	}
	return time.Time{}, false
}

// WhitelistListener 自定义Listener，用于过滤非白名单IP
type WhitelistListener struct {
	net.Listener
	whitelist *Whitelist
}

// NewWhitelistListener 创建新的白名单监听器
func NewWhitelistListener(listener net.Listener, whitelist *Whitelist) *WhitelistListener {
	return &WhitelistListener{
		Listener:  listener,
		whitelist: whitelist,
	}
}

// Accept 重写Accept方法，实现白名单过滤
func (wl *WhitelistListener) Accept() (net.Conn, error) {
	conn, err := wl.Listener.Accept()
	if err != nil {
		return nil, err
	}
	clientIP := conn.RemoteAddr().(*net.TCPAddr).IP.String()
	if !wl.whitelist.Contains(clientIP) {
		conn.Close()
		return nil, fmt.Errorf("IP %s not in whitelist", clientIP)
	}
	return conn, nil
}
